# Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved.
import hashlib
import subprocess
import uvicorn
from fastapi import FastAPI, File, UploadFile, Form,Request
import zipfile
import os
import shutil
import secrets
from pathlib import Path
import threading
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded
from slowapi.util import get_remote_address
from fastapi.responses import JSONResponse

from golden_seed_server.utils.corpus_manager.query_corpus import query_corpus
from golden_seed_server.utils.corpus_manager.delete_courpus import delete_corpus
from golden_seed_server.utils.corpus_manager.upload_courpus import upload_corpus, is_rag_busy
from golden_seed_server.utils.corpus_manager.change_doucument_to_para import change_document_to_para
from golden_seed_server.utils.corpus_manager.init_asset import init_asset
from golden_seed_server.content.config import config
from golden_seed_server.manager.user import UserManager

app = FastAPI(docs_url=None, redoc_url=None)
limiter = Limiter(key_func=get_remote_address)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

def calculate_uncompressed_size(zip_path):
    total_size = 0
    try:
        with zipfile.ZipFile(zip_path, 'r') as zip_file:
            for zip_info in zip_file.infolist():
                total_size += zip_info.file_size
    except:
        total_size=50 * 1024 * 1024+1
    return total_size

lock = threading.Lock()

@app.exception_handler(RateLimitExceeded)
async def too_many_requests(request, exc):
    return JSONResponse(content={"code": 429, "message": "服务正忙，请稍作休息再进行操作"}, status_code=429)

@app.post("/corpus")
async def upload_zip(
    request: Request = None,
    user_sub: str = Form("", min_length=0, max_length=20),
    user_passwd: str = Form("", min_length=0, max_length=20),
    method: str = Form("",  min_length=0, max_length=12),
    corpus_name: str = Form("", min_length=0, max_length=100),
    zip_file: UploadFile = File(None)
):
    passwd = hashlib.sha256(user_passwd.encode('utf-8')).hexdigest()
    user=UserManager.get_userinfo_by_user_sub(user_sub=user_sub)
    if user is None:
        return {"code":403,"message":"用户名不存在"}
    if user.passwd!=passwd:
        return {"code":403,"message":"用户密码错误"}
    if method == 'init_asset':
        init_asset(config['PG_HOST'], config['PG_PORT'], config['PG_USER'],
                   config['PG_PWD'], user_sub, user_passwd, 'BGE_MIXED_MODEL')
    elif method == 'up_corpus':
        lock.acquire()
        if is_rag_busy(config['PG_HOST'], config['PG_PORT'], config['PG_USER'], config['PG_PWD']):
            lock.release()
            return {"code": 201, "message": "等待之前任务完成中"}
        try:
            if not os.path.exists('./upload'):
                os.mkdir('./upload')
            os.chmod('./upload', 0o755)
            corpus_dir = os.path.join('./upload', user_sub, os.path.splitext(zip_file.filename)[0])
            if os.path.exists(corpus_dir):
                shutil.rmtree(corpus_dir)
            os.makedirs(corpus_dir)
            file_path = os.path.join('./upload', user_sub, zip_file.filename)
            f = open(file_path, "wb")
            data = await zip_file.read()
            f.write(data)
            f.close()
        except:
            lock.release()
        if calculate_uncompressed_size(file_path) > 10*1024 * 1024:
            lock.release()
            os.remove(file_path)
            return {"code": 400, "message": "文件太大"}
        try:
            with zipfile.ZipFile(file_path, 'r') as zip_ref:
                zip_ref.extractall(corpus_dir)
                print(f"ZIP文件已解压至{corpus_dir}")
        except zipfile.BadZipFile:
            lock.acquire()
            os.remove(file_path)
            return {"code": 400, "message": "上传的文件不是一个有效的ZIP文件"}
        finally:
            os.remove(file_path)
        para_dir = os.path.join('./upload', user_sub, secrets.token_hex(16))
        print("para_dir",para_dir)
        if os.path.exists(para_dir):
            if os.path.islink(para_dir):
                os.unlink(para_dir)
            shutil.rmtree(para_dir)
        os.mkdir(para_dir)
        corpus_name_list = change_document_to_para(corpus_dir, para_dir, config['CHUNK_SIZE'])
        try:
            print("qiegekais")
            corpus_name_list = change_document_to_para(corpus_dir, para_dir, config['CHUNK_SIZE'])
            print("qiegejieshu")
            for corpus_name in corpus_name_list:
                delete_corpus(
                    config['PG_HOST'],
                    config['PG_PORT'],
                    config['PG_USER'],
                    config['PG_PWD'],
                    user_sub, user_passwd, corpus_name)
            upload_thread = threading.Thread(
                target=upload_corpus,
                args=(config['PG_HOST'],
                      config['PG_PORT'],
                      config['PG_USER'],
                      config['PG_PWD'],
                      config['SSL_ENABLE'],
                      config['RAG_HOST'],
                      config['RAG_PORT'],
                      user_sub, user_passwd, para_dir))
            upload_thread.start()
            return {"code": 200, "message": "等待语料上传成功"}
        except Exception as e:
            return {"code": 400, "message": "语料上传失败"}
        finally:
            shutil.rmtree(corpus_dir)
            lock.release()
    elif method == 'del_corpus':
        try:
            delete_corpus(
                config['PG_HOST'],
                config['PG_PORT'],
                config['PG_USER'],
                config['PG_PWD'],
                user_sub, user_passwd, corpus_name)
            return {"code": 200, "message": "语料删除成功"}
        except:
            return {"code": 400, "message": "语料删除失败"}
    elif method == 'query_corpus':
        try:
            results = query_corpus(
                config['PG_HOST'],
                config['PG_PORT'],
                config['PG_USER'],
                config['PG_PWD'],
                user_sub, user_passwd, corpus_name)
            return {"code": 200, "message": results}
        except:
            return {"code": 400, "message": "语料查询失败"}
    return {"code": 400, "message": "不支持当前模式"}

if __name__ == "__main__":
    uvicorn.run(
        app,
        host=config["UVICORN_HOST"],
        port=int(config["UVICORN_PORT"])
    )
